# -*- coding: utf-8 -*-
# time: 2025/4/17 11:11
# file: video_toText_hf.py
# author: hanson
"""
案例代码：BLIP-2 实现视频转文本

"""
import cv2
import torch
from decord import VideoReader, cpu
from transformers import Blip2Processor, Blip2ForConditionalGeneration
from typing import List


def extract_keyframes(video_path: str, interval: int = 10) -> List[torch.Tensor]:
    """提取视频关键帧（每隔interval秒取一帧）"""
    vr = VideoReader(video_path, ctx=cpu(0))
    fps = vr.get_avg_fps()
    frame_indices = range(0, len(vr), int(fps * interval))
    frames = [vr[i].asnumpy() for i in frame_indices]
    return [torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) for frame in frames]


def video_to_text(video_path: str, model_name: str = "Salesforce/blip2-opt-2.7b") -> str:
    # 初始化模型（自动下载，约5GB）
    processor = Blip2Processor.from_pretrained(model_name)
    model = Blip2ForConditionalGeneration.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto"
    )

    # 提取关键帧并生成描述
    frames = extract_keyframes(video_path)
    descriptions = []
    for frame in frames:
        inputs = processor(images=frame, return_tensors="pt").to("cuda", torch.float16)
        out = model.generate(**inputs, max_new_tokens=50)
        desc = processor.decode(out[0], skip_special_tokens=True)
        descriptions.append(desc)

    return " | ".join(descriptions)


# 测试示例
if __name__ == "__main__":
    video_path = "test.mp4"  # 替换为你的视频路径
    print(video_to_text(video_path))